"""
核反应堆子通道分析程序的Numba优化函数
"""

import numpy as np
from numba import jit, njit, prange
import numba as nb
from constants.constants import (
    GRAVITY, CFL_NUMBER, TURBULENT_MIXING_COEFF,
    INTERFACE_DRAG_COEFF, SUBCOOLED_BOILING_COEFF
)

@njit(parallel=True)
def compute_grid_points_numba(n_axial, n_channels, void_fraction, prev_void_fraction, 
                            vapor_velocity, liquid_velocity, prev_liquid_velocity,
                            prev_vapor_velocity, pressure, temperature, time_step, dz):
    """使用Numba优化的网格点计算"""
    new_void = np.zeros_like(void_fraction)
    new_vl = np.zeros_like(liquid_velocity)
    new_vv = np.zeros_like(vapor_velocity)
    
    for i in prange(1, n_axial):
        for j in prange(n_channels):
            # 计算空间导数项
            if i > 0:
                d_alpha_dx = (void_fraction[i,j] - void_fraction[i-1,j]) / dz
                d_vl_dx = (liquid_velocity[i,j] - liquid_velocity[i-1,j]) / dz
                d_vv_dx = (vapor_velocity[i,j] - vapor_velocity[i-1,j]) / dz
                dp_dx = (pressure[i,j] - pressure[i-1,j]) / dz
            else:
                d_alpha_dx = 0
                d_vl_dx = 0
                d_vv_dx = 0
                dp_dx = 0
            
            # 更新值
            new_void[i,j] = prev_void_fraction[i,j] - time_step * (
                vapor_velocity[i,j] * d_alpha_dx
            )
            
            new_vl[i,j] = prev_liquid_velocity[i,j] - time_step * (
                liquid_velocity[i,j] * d_vl_dx +
                dp_dx +
                GRAVITY
            )
            
            new_vv[i,j] = prev_vapor_velocity[i,j] - time_step * (
                vapor_velocity[i,j] * d_vv_dx +
                dp_dx +
                GRAVITY
            )
            
            # 限制物理范围
            new_void[i,j] = min(max(new_void[i,j], 0.0), 0.9999)
            new_vl[i,j] = max(new_vl[i,j], 0.0)
            new_vv[i,j] = max(new_vv[i,j], 0.0)
            
    return new_void, new_vl, new_vv

@njit(parallel=True)
def compute_energy_points_numba(n_axial, n_channels, temperature, prev_temperature,
                              liquid_velocity, pressure, time_step, dz):
    """使用Numba优化的能量点计算"""
    new_temp = np.zeros_like(temperature)
    
    for i in prange(1, n_axial):
        for j in prange(n_channels):
            # 计算空间导��项
            if i > 0:
                d_T_dx = (temperature[i,j] - temperature[i-1,j]) / dz
            else:
                d_T_dx = 0
            
            # 更新温度
            new_temp[i,j] = prev_temperature[i,j] - time_step * (
                liquid_velocity[i,j] * d_T_dx
            )
            
    return new_temp

@njit
def calculate_wall_friction_numba(velocity, diameter, density, viscosity):
    """使用Numba优化的壁面摩擦力计算"""
    # 计算Reynolds数
    Re = density * abs(velocity) * diameter / viscosity
    
    # 计算摩擦系数
    if Re < 2300:  # 层流
        f = 64 / Re
    else:  # 湍流
        f = 0.316 * Re**(-0.25)
        
    # 计算壁面摩擦力
    tau_w = 0.5 * f * density * velocity * abs(velocity)
    
    return tau_w

@njit(parallel=True)
def calculate_residual_numba(n_axial, n_channels, liquid_velocity, liquid_enthalpy, dz):
    """使用Numba优化的残差计算"""
    mass_residual = np.zeros((n_axial, n_channels))
    energy_residual = np.zeros((n_axial, n_channels))
    
    for i in prange(1, n_axial):
        for j in prange(n_channels):
            # 轴向质量流量差
            dm_dz = (liquid_velocity[i,j] - liquid_velocity[i-1,j]) / dz
            mass_residual[i,j] = dm_dz
            
            # 轴向能量���运
            dh_dz = (liquid_enthalpy[i,j] - liquid_enthalpy[i-1,j]) / dz
            energy_residual[i,j] = dh_dz
            
    return max(np.max(np.abs(mass_residual)), np.max(np.abs(energy_residual)))

@njit
def calculate_turbulent_mixing_numba(alpha_i, alpha_j, rho, dy):
    """计算湍流混合项"""
    L_tur = dy  # 湍流混合长度
    alpha_eq_i = 0.5  # 平衡空泡分数
    
    return TURBULENT_MIXING_COEFF * dy * L_tur * rho * (alpha_j - alpha_i - (alpha_j - alpha_eq_i))

@njit
def calculate_interfacial_force_numba(alpha_l, rho_l, u_l, u_v, v_l, v_v):
    """计算界面力"""
    d_b = 0.001  # 气泡直径
    
    # 计算相对速度
    u_r = u_v - u_l
    v_r = v_v - v_l
    V_r = np.sqrt(u_r**2 + v_r**2)
    
    # 计算界面面积密度
    a_i = 6 * (1-alpha_l) / d_b
    
    return 0.5 * INTERFACE_DRAG_COEFF * rho_l * V_r * V_r * a_i

@njit
def calculate_wall_heat_flux_numba(T_w, T_f, h):
    """计算壁面热流密度"""
    return h * (T_w - T_f)

@njit
def calculate_saturation_temperature_numba(P):
    """计算饱和温度"""
    return 373.15 + 0.0001 * (P - 1.013e5)  # 简化的饱和温度计算

@njit
def calculate_latent_heat_numba(P):
    """计算汽化潜热"""
    return 2.257e6  # 常压下水的汽化潜热

@njit
def calculate_cfl_timestep(u_l, u_v, v_l, v_v, dx, dy, dz):
    """计算基于CFL条件的时间步长"""
    # 计算最大速度
    max_velocity = max(
        np.max(np.abs(u_l)),
        np.max(np.abs(u_v)),
        np.max(np.abs(v_l)),
        np.max(np.abs(v_v)))
    
    # CFL条件，取最小网格间距
    min_dx = min(dx, dy, dz)
    
    # 计算时间步长
    dt = CFL_NUMBER * min_dx / (max_velocity + 1e-10)
    
    return dt

@njit
def check_numerical_stability(alpha_l, alpha_v, u_l, u_v, v_l, v_v, P, T_l, T_v,
                            rho_l, rho_v, h_l, h_v, dt_old):
    """检查数值稳定性并调整时间步长"""
    # 1. 检查物理量的变化率
    max_alpha_change = max(
        np.max(np.abs(np.diff(alpha_l))),
        np.max(np.abs(np.diff(alpha_v)))
    )
    
    max_velocity_change = max(
        np.max(np.abs(np.diff(u_l))),
        np.max(np.abs(np.diff(u_v))),
        np.max(np.abs(np.diff(v_l))),
        np.max(np.abs(np.diff(v_v)))
    )
    
    max_pressure_change = np.max(np.abs(np.diff(P))) / np.mean(P)
    max_temp_change = max(
        np.max(np.abs(np.diff(T_l))) / np.mean(T_l),
        np.max(np.abs(np.diff(T_v))) / np.mean(T_v)
    )
    
    # 2. 设置阈值
    alpha_threshold = 0.1
    velocity_threshold = 0.2
    pressure_threshold = 0.1
    temp_threshold = 0.1
    
    # 3. 根据变化率调整时间步长
    dt_factor = 1.0
    
    if max_alpha_change > alpha_threshold:
        dt_factor = min(dt_factor, 0.5)
    if max_velocity_change > velocity_threshold:
        dt_factor = min(dt_factor, 0.5)
    if max_pressure_change > pressure_threshold:
        dt_factor = min(dt_factor, 0.5)
    if max_temp_change > temp_threshold:
        dt_factor = min(dt_factor, 0.5)
        
    # 如果所有变化率都很小，可以适当增加时间步长
    if (max_alpha_change < alpha_threshold/2 and
        max_velocity_change < velocity_threshold/2 and
        max_pressure_change < pressure_threshold/2 and
        max_temp_change < temp_threshold/2):
        dt_factor = min(1.2, 1.0/dt_factor)
        
    # 4. 检查物理约束
    if (np.any(alpha_l < 0) or np.any(alpha_l > 1) or
        np.any(P < 0) or
        np.any(T_l < 273) or np.any(T_v < 273) or
        np.any(rho_l < 0) or np.any(rho_v < 0)):
        dt_factor = 0.1  # 严重违反物理约束，大幅减小时间步长
        
    # 5. 限制时间步长的变化范围
    dt_new = dt_old * dt_factor
    dt_new = max(1e-6, min(dt_new, 1e-2))  # 限制在合理范围内
    
    return dt_new, dt_factor < 1.0

@njit
def calculate_fuel_heat_generation_numba(i, n_axial, dz):
    """计算燃料热生成 - Numba优化版本"""
    q_max = 1e6  # 最大线功率密度 (W/m)
    z = i * dz
    H = n_axial * dz
    # 使用余弦分布模拟轴向功率分布
    q = q_max * np.cos(np.pi * z / H)
    return q

@njit
def calculate_radial_heat_conduction_numba(T_c, T_f, dr):
    """计算径向热传导 - Numba优化版本"""
    k = 30.0  # 假设的导热系数 (W/m·K)
    # 简化的径向热传导计算
    q_rad = k * (T_c - T_f) / dr
    return q_rad
